
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

original_img = Image.open('input.png').convert('L').resize((512, 512))
original_array = np.array(original_img)

dft = np.fft.fft2(original_array)
dft_shift = np.fft.fftshift(dft)
magnitude_spectrum = 20 * np.log(np.abs(dft_shift))

watermark = Image.open('custom_watermark.png').convert('L').resize((128, 128))
watermark_array = np.where(np.array(watermark) < 128, 1,0) * 50000  

h, w = dft_shift.shape
dft_shift[:128, :128] += watermark_array
dft_shift[-128:, -128:] += watermark_array[::-1, ::-1]
watermarked_spectrum = 30 * np.log(np.abs(dft_shift))

dft_ishift = np.fft.ifftshift(dft_shift)
img_back = np.fft.ifft2(dft_ishift)
img_back = np.abs(img_back).astype(np.uint8)

plt.figure(figsize=(15, 10))

plt.subplot(2, 2, 1)
plt.imshow(original_array, cmap='gray')
plt.title('Original Image')

plt.subplot(2, 2, 2)
plt.imshow(magnitude_spectrum, cmap='gray')
plt.title('Frequency Domain (Original)')

plt.subplot(2, 2, 3)
plt.imshow(watermarked_spectrum, cmap='gray')
plt.title('Watermarked Frequency Domain')

plt.subplot(2, 2, 4)
plt.imshow(img_back, cmap='gray')
plt.title('Reconstructed Image')

plt.tight_layout()
plt.show()